/*
 *	Copyright © 2008 University of Houston
 *	All rights reserved.
 */

#include <pcap.h>
#include <stdio.h>
#include <string>
#include <stdlib.h>
#include <ctype.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <iostream>
#include "listener.h"
#include "session.h"
#include "logger.h"
#include "analyzer.h"
#include "filter.h"
#include "util.h"

using namespace std;


Listener::Listener(string device, string portNum, string hosts, string file,
				   int filt, int alg)
{
	dev 			= device;
	logFilename 	= file;
	hostString		= hosts;
	filterID		= filt;
	algID			= alg;

	if (!logger.init(logFilename))
	{
		FATAL("Could not initialize logger")
		return;
	}

	sessionCount = logger.getLastSessionNumber();
	INFO("Determined next session number is " << sessionCount)

	// Ugly code to build a filter expression
	filter_exp = "";
	vector<string> tcpHosts = stringSplit(hosts, ",");

	if (tcpHosts.size() > 0)
	{
		filter_exp += "(host ";
		for (u_int hostNum = 0; hostNum < tcpHosts.size(); hostNum++)
		{
			filter_exp += tcpHosts[hostNum];
			if (hostNum != tcpHosts.size()-1)
				filter_exp += " or ";
		}
		filter_exp += ") and ";
	}

	if (portNum == "")
		filter_exp += "tcp port 22";
	else
		filter_exp += "tcp port " + portNum;

	INFO("Filter Expression: " << filter_exp)

	if (!initializePcap())
		ERROR("Could not initialize pcap!")
};

Listener::~Listener()
{

	// Save all of the sessions that have not closed to file.
	while (sessionList.size() > 0)
		closeSession((sessionList[0])->getSessionID());

	// Cleanup pcap
	pcap_freecode(&pcapData.filterProgram);
	pcap_close(pcapData.handle);
}

void Listener::run()
{
	/* now we can set our callback function */
	pcap_loop(pcapData.handle, -1, got_packet, (u_char *)this);
}


void got_packet(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{
	Listener *lis		 				= (Listener*)args;
	session_list *sessionList 			= &(lis->sessionList);
	session_list::iterator				sessIter;
	timedHeader 						pktHeader;

	// Parse the packet information.  If the packet is bad, do not
	// proceed.
	if (!pktHeader.formatPacket(packet))
		return;

	// Assign the packet to a session.
	bool foundSession = false;

	for (sessIter = sessionList->begin();
			sessIter != sessionList->end();
			sessIter++)
	{

		// If we have found the session object that the packet blongs to...
		if ((*sessIter)->belongs(pktHeader))
		{
			// Make sure not to create a new session below
			foundSession = true;

			// Add the packet to the session
			(*sessIter)->addPacket(pktHeader);

			// Check for a fin-ack sequence
			if ((*sessIter)->isConnectionClosed())
				lis->closeSession((*sessIter)->getSessionID());
			break;
		}
	}

	// There is not a session object that this packet belongs to.
	// Create a new session object, and a corresponding analyzer object.
	if (!foundSession)
		lis->newSession(pktHeader);

	INFO(lis->logger.formatPacketForLogging(pktHeader))
};

bool Listener::setFilter(string s)
{
	// Get our filter expression into a usable buffer
	char *buffer = new char[s.size()+1];
	strcpy(buffer, s.c_str());

	// Compile the expression
	if (pcap_compile(pcapData.handle, &pcapData.filterProgram, buffer, 0, pcapData.net) == -1)
	{
		ERROR("Couldn't parse filter " << filter_exp << ": " 
				<< pcap_geterr(pcapData.handle))
		return false;
	}

	// Set the filter
	if (pcap_setfilter(pcapData.handle, &pcapData.filterProgram) == -1)
	{
		ERROR("Couldn't install filter " << filter_exp << " "
				<< pcap_geterr(pcapData.handle))
		return false;
	}

	// Cleanup
	delete [] buffer;

	return true;
};

string Listener::getDevice()
{

	if (dev.length() == 0)
	{
		dev = pcap_lookupdev(pcapData.errbuf);

		if (dev.length() == 0)
			ERROR("Could not open device " << dev << pcapData.errbuf)
	}

	return dev;
};

pcap_t* Listener::openDevice()
{
	// We don't care about the packet data itself.  Ignore the data.
	const int maxPacketData = SNAP_LEN;

	// Amount of time before we timeout on listening (in ms)
	const int timeoutDelay = 1000;

	// Put the device in promiscuous mode?
	const int promiscMode = 0;

	// Don't care what this does
	bpf_u_int32 	mask;

	// Actually open the device
	pcapData.handle = pcap_open_live(dev.c_str(),
									 maxPacketData,
									 promiscMode,
									 timeoutDelay,
									 pcapData.errbuf);

	if (pcapData.handle == NULL)
		ERROR("Couldn't open device " << dev)

	else if (pcap_lookupnet(dev.c_str(), &pcapData.net, &mask, pcapData.errbuf) == -1)
	{
		ERROR("Couldn't get netmask for device " << dev)
		pcapData.net = 0;
		return NULL;
	}

	return pcapData.handle;
};

bool Listener::initializePcap()
{
	// Set up our capture device.
	if (dev.length() == 0)
		dev = getDevice();

	if (openDevice() == NULL)
		return false;

	if (!setFilter(filter_exp))
		return false;

	return true;
};

const session* Listener::operator[](int index)
{
	return sessionList[index];
};

int Listener::addSession(session *s)
{
	// Give it an ID number
	s->setSessionID(sessionCount);

	// Keep track of it
	sessionList.push_back(s);

	// Return the session number that was assigned.
	return sessionCount++;
}

int Listener::newSession(timedHeader h)
{
	// Create the new object
	session* s 		= new session();

	// Set up the Analyzer for this session
	sessionAnalyzer* a		= new sessionAnalyzer();

	// Initialize Analyzer to print results to screen
	a->initialize(s, filter::getFilter(filterID), algorithm::getAlg(algID));

	s->setAnalyzer(a);

	s->addPacket(h);

	// Keep track of it
	return addSession(s);
}

int Listener::closeSession(int sessionID)
{
	session_list::iterator sessIter;

	for (sessIter = sessionList.begin();
			sessIter != sessionList.end();
			sessIter++)
	{
		if ((*sessIter)->getSessionID() == sessionID)
		{
			// Log the entire session to file
			logger.logSession((*sessIter));

			// Free up the memory
			delete(*sessIter);

			// Remove the session from tracking
			sessionList.erase(sessIter);

			return sessionID;
		}
	}

	return -1;
}

